# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com

from quantlib import *
import MySQLLoading.loadData
from datetime import datetime, timedelta
from copy import deepcopy
from PerformanceMetrics import *


def convertToQLDate(inputDate):
    return Date(inputDate.day, inputDate.month, inputDate.year)

def convertFromQLDate(inputDate):
    return datetime.date(datetime(inputDate.year(), inputDate.month(), inputDate.dayOfMonth()))

def makeDateRateHolder(ticker, dateLower, dateHigher, dates, table = 'historicalDly'):
    data = MySQLLoading.loadData.LoadSQLDateRate(table, ticker, dateLower, dateHigher, True if table == 'historicalDly' else False)
    for i in range(len(data[0])):
        dates.append(data[0][i]) 
    drh = DateRateHolder()
    for idx in range(len(data[0])):
        drh.setElement(convertToQLDate(data[0][idx]),data[1][idx])
    return drh

def convertDateRateHolder(drh):
    length = drh.size()
    x = []
    y = []
    for idx in range(length):
        tempDate = drh.getElementDate(idx)
        x.append(convertFromQLDate(tempDate))
        y.append(drh.getElementRate(idx))
    return [x,y]


class tickerFactory:

    def __init__(self, tickers, haveDividends=[], dateLower=datetime.date(datetime(1901,1,1)), dateHigher=datetime.date(datetime.today())):
        self.items = {}
        self.availableDates = []
        for ticker in tickers:
            tempDates = []
            self.items[ticker] = makeDateRateHolder(ticker, dateLower, dateHigher, tempDates)
            self.availableDates = [d for d in tempDates if d in self.availableDates] if ticker != tickers[0] else tempDates
            if ticker in haveDividends:
                self.items['dividends_' + ticker] = makeDateRateHolder(ticker, dateLower, dateHigher, tempDates, 'historicaldividends')


class backtestEngine:

    backtestStrategies = {'long': LongStrategy}

    def __init__(self, initData):
        
        self.portfolio = Portfolio()
        self.availableDates = initData.availableDates
        self.availableDates.sort()
        self.portfolio.initialisePortfolio(convertToQLDate(self.availableDates[0]), 1000000)
        for drhName, drhData in initData.items.iteritems():
            if not 'dividends_' in drhName:
                if 'dividends_' + drhName in initData.items.keys():
                    newAsset = AssetData(drhData, initData.items['dividends_' + drhName])
                else:
                    newAsset = AssetData(drhData)
                self.portfolio.addAsset(drhName, newAsset)

    def addResults(self, name, resultDict):
        for itemName, itemVal in resultDict.iteritems():
            self.backtestResults[name + '.' + itemName] = itemVal
      
    def backtest(self, strategyType, firstStartDate, startDatesDeltaInDays, lengthInDays, currency, *args):
        # Sort out the start dates first:
        startDates = []
        self.backtestResults = {}
        availableDates = self.availableDates
        finalAvailableDate = availableDates[-1]
        for date in availableDates:
            #   Find the first eligible start date
            if date + timedelta(lengthInDays) <= finalAvailableDate and firstStartDate <= date:
                startDates.append(convertToQLDate(date))
                lastDate = date
                break

        if len(startDates) == 0:
            raise Exception('No eligible start dates')
        
        for date in availableDates:
            #   Now continue adding start dates until we run out of space
            if date + timedelta(lengthInDays) > finalAvailableDate:
                break
            if date >= lastDate + timedelta(startDatesDeltaInDays):
                startDates.append(convertToQLDate(date))
                lastDate = date

        qlStartDates = DateVector(len(startDates))
        for idx in range(len(startDates)):
            qlStartDates[idx] = startDates[idx]
        strategy = backtestEngine.backtestStrategies[strategyType](*args)
        eng = TradingEngine(self.portfolio, strategy)
        performanceVec = eng.performBacktest(qlStartDates, lengthInDays)
        results = {}
        positions = []
        for performance in performanceVec:
            obs = performance.getObservations()
            positions.append(convertDateRateHolder(obs))

        #   get investment ratios
        self.addResults('sharpe', sharpeRatio(positions))
        self.addResults('sortino', sortinoRatio(positions))

        #   get drawdown stats
        self.addResults('drawdown', drawdown(positions))

    def printResults(self):
        if not self.backtestResults is None:
            myKeys = self.backtestResults.keys()
            myKeys.sort()
            for key in myKeys:
                print(key + ' = ' + str(self.backtestResults[key]))
        